# -*- coding: utf8 -*-

#
# Some code is taken from http://www.cnblogs.com/x1957/archive/2012/06/02/2531503.html
#

# -*- coding: utf8 -*-

import struct

import numpy as np


class ImageFile(object):
    """
    The image file object to read the images.
    """

    def __init__(self, filename):
        super(ImageFile, self).__init__()
        self.filename = filename
        with open(filename, 'rb') as src:
            read_size = struct.calcsize('>IIII')
            buf = src.read(read_size)
            self.magic, self.numImages, self.numRows, self.numColumns = struct.unpack_from('>IIII', buf, 0)

    def images(self):
        """
        Read images from the binary file.
        :return: generator
        """
        head_size = struct.calcsize('>IIII')
        image_size = struct.calcsize('>%dB' % (self.numRows * self.numColumns))
        with open(self.filename, 'rb') as src:
            buf = src.read()
            for index in range(self.numImages):
                im = struct.unpack_from('>784B', buf, head_size + index * image_size)

                im = np.array(im)
                im = im.reshape(28, 28)

                yield im

class LabelFile(object):
    """
    Label file structure object.
    """

    def __init__(self, filename):
        super(LabelFile, self).__init__()
        self.filename = filename
        with open(self.filename, 'rb') as src:
            read_size = struct.calcsize('>II')
            buf = src.read(read_size)
            self.magic, self.numLabels = struct.unpack_from('>II', buf)

    def labels(self):
        """
        Iterate all the labels.
        :return: generator.l
        """
        head_size = struct.calcsize('>II')
        label_size = struct.calcsize('>1B')
        with open(self.filename, 'rb') as src:
            buf = src.read()
            for index in range(self.numLabels):
                label = struct.unpack_from('>1B', buf, head_size + index * label_size)[0]
                yield label
